from machine import SPI, Pin
import utime
import st7789
import urandom

# 按钮初始化
UP = Pin(1)
DOWN = Pin(0)
LEFT = Pin(2)
RIGHT = Pin(14)
FIRE = Pin(13)  # 射击按钮

UP.init(Pin.IN)
DOWN.init(Pin.IN)
LEFT.init(Pin.IN)
RIGHT.init(Pin.IN)
FIRE.init(Pin.IN)

UP.pull(Pin.PULL_DOWN)
DOWN.pull(Pin.PULL_DOWN)
LEFT.pull(Pin.PULL_DOWN)
RIGHT.pull(Pin.PULL_DOWN)
FIRE.pull(Pin.PULL_DOWN)

# 显示初始化
spi = SPI(0, baudrate=40000000, polarity=1, phase=0, bits=8, endia=0, sck=Pin(6), mosi=Pin(8))
display = st7789.ST7789(spi, 240, 240, reset=Pin(11,func=Pin.GPIO, dir=Pin.OUT), dc=Pin(7,func=Pin.GPIO, dir=Pin.OUT))
display.init()

# 定义窗口的大小
windowsWidth = 240
windowsHeight = 240

# 定义基础单位大小
cellSize = 10 

# 定义地图大小
mapWidth = int(windowsWidth / cellSize)
mapHeight = int(windowsHeight / cellSize)

# 颜色定义
MapColor = st7789.BLACK
PlayerColor = st7789.GREEN
EnemyColor = st7789.RED
BulletColor = st7789.YELLOW
ObstacleColor = st7789.WHITE  # 使用白色作为障碍物颜色
HighlightColor = st7789.BLUE  # 新增高亮颜色

# 游戏状态
PLAYING = 0
GAME_OVER = 1
START_SCREEN = 2  # 新增开始界面状态
game_state = START_SCREEN  # 初始状态为开始界面

# 分数
player_score = 0
enemy_score = 0

# 坦克类
class Tank:
    def __init__(self, x, y, direction, color, is_player=False):
        self.x = x
        self.y = y
        self.direction = direction  # 0:上, 1:右, 2:下, 3:左
        self.color = color
        self.is_player = is_player
        self.alive = True
        self.fire_cooldown = 0
        self.prev_x = x
        self.prev_y = y
        
    def draw(self):
        if not self.alive:
            return
            
        # 绘制坦克主体
        DrawCell(self.x, self.y, self.color)
        
    def move(self, dx, dy, obstacles, all_tanks):
        # 保存移动前的位置
        self.prev_x = self.x
        self.prev_y = self.y
        
        new_x = self.x + dx
        new_y = self.y + dy
        
        # 检查边界
        if new_x < 0 or new_x >= mapWidth or new_y < 0 or new_y >= mapHeight:
            return False
            
        # 检查障碍物
        if (new_x, new_y) in obstacles:
            return False
            
        # 检查是否与其他坦克重叠（包括玩家和敌人）
        for tank in all_tanks:
            if tank != self and tank.alive and new_x == tank.x and new_y == tank.y:
                return False
                
        # 清除当前位置
        if (self.x, self.y) in obstacles:
            DrawCell(self.x, self.y, ObstacleColor)
        else:
            DrawCell(self.x, self.y, MapColor)
                
        self.x = new_x
        self.y = new_y
        
        # 绘制新位置
        self.draw()
        
        return True
        
    def fire(self, bullets, obstacles):
        if self.fire_cooldown > 0:
            return False
            
        # 根据方向确定子弹初始位置
        if self.direction == 0:  # 上
            bullet_x, bullet_y = self.x, self.y - 1
        elif self.direction == 1:  # 右
            bullet_x, bullet_y = self.x + 1, self.y
        elif self.direction == 2:  # 下
            bullet_x, bullet_y = self.x, self.y + 1
        else:  # 左
            bullet_x, bullet_y = self.x - 1, self.y
        
        # 检查子弹初始位置是否有障碍物
        if IsObstacle(bullet_x, bullet_y, obstacles):
            return False  # 不能发射子弹
            
        bullets.append([bullet_x, bullet_y, self.direction, self.is_player])  # 简化子弹数据结构
        self.fire_cooldown = 10
        return True
        
    def update(self):
        if self.fire_cooldown > 0:
            self.fire_cooldown -= 1

# 绘制单元格
def DrawCell(x, y, color):
    display.fill_rect(x * cellSize, y * cellSize, cellSize, cellSize, color)

# 绘制地图 - 移除了边界绘制
def DrawMap(obstacles):
    # 只绘制障碍物，不绘制边界
    for x, y in obstacles:
        DrawCell(x, y, ObstacleColor)

# 绘制分数
def DrawScore():
    score_text = "P:" + str(player_score) + " E:" + str(enemy_score)
    display.draw_string(5, 5, score_text, color=PlayerColor, bg=MapColor, size=2)
    
# 绘制难度等级
def DrawDifficultyLevel(enemy_count):
    # 计算难度等级 (1-3)
    difficulty = min(enemy_count, 3)
    # 在右上角显示难度等级
    difficulty_text = "Lv." + str(difficulty)
    # 计算文本宽度以便右对齐 (假设每个字符大约12像素宽)
    text_width = len(difficulty_text) * 12
    x_pos = windowsWidth - text_width - 5  # 右边留5像素边距
    display.draw_string(x_pos, 5, difficulty_text, color=EnemyColor, bg=MapColor, size=2)
    
# 开始界面
def StartScreen():
    display.fill(MapColor)
    display.draw_string(40, 80, "TANK BATTLE", color=PlayerColor, bg=MapColor, size=3)
    display.draw_string(30, 130, "Press any key", color=BulletColor, bg=MapColor, size=2)
    display.draw_string(40, 160, "to start", color=BulletColor, bg=MapColor, size=2)

# 游戏结束画面 - 修改为有两个选项
def GameOver(selected_option):
    display.fill(MapColor)
    
    #胜负判断逻辑：根据分数
    if player_score > enemy_score:
        display.draw_string(40, 40, "YOU WIN!", color=PlayerColor, bg=MapColor, size=3)
    elif player_score < enemy_score:
        display.draw_string(40, 40, "YOU LOSE!", color=EnemyColor, bg=MapColor, size=3)
    else:
        display.draw_string(40, 40, "DRAW!", color=BulletColor, bg=MapColor, size=3)
    
    score_text = "Score: " + str(player_score) + "-" + str(enemy_score)
    display.draw_string(30, 80, score_text, color=BulletColor, bg=MapColor, size=2)
    
    # 添加两个选项
    if selected_option == 0:
        # 继续游戏选项高亮
        display.draw_string(50, 120, "Continue", color=HighlightColor, bg=MapColor, size=2)
        display.draw_string(50, 150, "Main Menu", color=PlayerColor, bg=MapColor, size=2)
    else:
        # 主菜单选项高亮
        display.draw_string(50, 120, "Continue", color=PlayerColor, bg=MapColor, size=2)
        display.draw_string(50, 150, "Main Menu", color=HighlightColor, bg=MapColor, size=2)

# 简单AI
def SimpleAI(tank, player_tank, obstacles, bullets, all_tanks):
    # 生成一个随机数决定AI行为
    action = urandom.randint(0, 100)
    
    if action < 15:  # 15%的概率移动
        # 计算与玩家的方向差异
        dx = player_tank.x - tank.x
        dy = player_tank.y - tank.y
        
        # 决定移动方向
        if abs(dx) > abs(dy):
            if dx > 0:
                tank.direction = 1  # 右
                tank.move(1, 0, obstacles, all_tanks)
            else:
                tank.direction = 3  # 左
                tank.move(-1, 0, obstacles, all_tanks)
        else:
            if dy > 0:
                tank.direction = 2  # 下
                tank.move(0, 1, obstacles, all_tanks)
            else:
                tank.direction = 0  # 上
                tank.move(0, -1, obstacles, all_tanks)
    elif action < 30: 
        # 尝试射击
        tank.fire(bullets, obstacles)
    # 剩下什么也不做

# 检查位置是否有障碍物
def IsObstacle(x, y, obstacles):
    # 检查边界
    if x < 0 or x >= mapWidth or y < 0 or y >= mapHeight:
        return True
    
    # 检查障碍物
    if (x, y) in obstacles:
        return True
        
    return False

# 检查子弹碰撞
def CheckBulletCollision(bullets):
    # 创建一个字典来记录每个位置的子弹数量
    bullet_positions = {}
    
    # 记录每个位置的子弹索引
    for i, bullet in enumerate(bullets):
        pos = (bullet[0], bullet[1])
        if pos not in bullet_positions:
            bullet_positions[pos] = []
        bullet_positions[pos].append(i)
    
    # 找出有多个子弹的位置
    bullets_to_remove = []
    for pos, indices in bullet_positions.items():
        if len(indices) > 1:  # 如果同一位置有多个子弹
            bullets_to_remove.extend(indices)
    
    # 从后往前移除子弹，避免索引变化问题
    bullets_to_remove.sort(reverse=True)
    for index in bullets_to_remove:
        if index < len(bullets):
            bullets.pop(index)

# 检查是否有任意按键按下
def AnyKeyPressed():
    return UP.value() == 1 or DOWN.value() == 1 or LEFT.value() == 1 or RIGHT.value() == 1 or FIRE.value() == 1

# 检查是否有方向键按下
def DirectionKeyPressed():
    return UP.value() == 1 or DOWN.value() == 1 or LEFT.value() == 1 or RIGHT.value() == 1

# 主游戏函数
def main():
    global game_state, player_score, enemy_score
    
    # 游戏难度级别
    difficulty_level = 1  # 初始难度级别
    
    # 创建障碍物
    obstacles = []
    for i in range(40):  # 从40个
        x = urandom.randint(2, mapWidth - 3)
        y = urandom.randint(2, mapHeight - 3)
        obstacles.append((x, y))
    
    # 初始绘制整个地图
    display.fill(MapColor)
    DrawMap(obstacles)
    
    # 创建玩家坦克
    player_tank = Tank(2, 2, 1, PlayerColor, True)
    
    # 敌人坦克列表
    enemy_tanks = []
    
    # 根据难度级别添加敌人坦克
    if difficulty_level >= 1:  # 初始至少有一个敌人
        enemy_tanks.append(Tank(mapWidth - 3, mapHeight - 3, 3, EnemyColor))
    
    if difficulty_level >= 2:  # 难度级别2增加一个敌人
        enemy_tanks.append(Tank(2, mapHeight - 3, 0, EnemyColor))
    
    if difficulty_level >= 3:  # 难度级别3再增加一个敌人
        enemy_tanks.append(Tank(mapWidth - 3, 2, 2, EnemyColor))
    
    # 子弹列表 [x, y, direction, is_player]
    bullets = []
    prev_bullets = []  # 存储上一帧的子弹位置
    
    # 游戏结束界面选项 (0:继续游戏, 1:返回主页面)
    game_over_option = 0
    
    # 初始绘制分数和难度等级（只执行一次）
    DrawScore()
    DrawDifficultyLevel(len(enemy_tanks))
    
    # 初始绘制开始界面
    StartScreen()
    
    # 游戏循环
    while True:
        if game_state == START_SCREEN:
            # 检查是否有任意按键按下
            if AnyKeyPressed():
                game_state = PLAYING
                # 重置游戏状态
                player_score = 0
                enemy_score = 0
                difficulty_level = 1  # 重置难度级别
                # 重新创建障碍物
                obstacles = []
                for i in range(40):
                    x = urandom.randint(2, mapWidth - 3)
                    y = urandom.randint(2, mapHeight - 3)
                    obstacles.append((x, y))
                
                # 重新创建玩家坦克
                player_tank = Tank(2, 2, 1, PlayerColor, True)
                
                # 重新创建敌人坦克
                enemy_tanks = []
                enemy_tanks.append(Tank(mapWidth - 3, mapHeight - 3, 3, EnemyColor))
                
                # 清空子弹
                bullets = []
                prev_bullets = []
                
                # 重新绘制地图和界面元素
                display.fill(MapColor)
                DrawMap(obstacles)
                DrawScore()
                DrawDifficultyLevel(len(enemy_tanks))
                # 等待按键释放
                while AnyKeyPressed():
                    utime.sleep_ms(50)
        
        elif game_state == PLAYING:
            # 创建所有坦克列表，用于碰撞检测
            all_tanks = [player_tank] + enemy_tanks
            
            # 清除上一帧的子弹位置
            for x, y, _, _ in prev_bullets:
                if (x, y) in obstacles:
                    DrawCell(x, y, ObstacleColor)
                else:
                    DrawCell(x, y, MapColor)
            
            # 保存当前子弹位置作为下一帧的上一帧位置
            prev_bullets = bullets.copy()
            
            # 处理玩家输入
            if UP.value() == 1:
                player_tank.direction = 0
                player_tank.move(0, -1, obstacles, all_tanks)
            elif DOWN.value() == 1:
                player_tank.direction = 2
                player_tank.move(0, 1, obstacles, all_tanks)
            elif LEFT.value() == 1:
                player_tank.direction = 3
                player_tank.move(-1, 0, obstacles, all_tanks)
            elif RIGHT.value() == 1:
                player_tank.direction = 1
                player_tank.move(1, 0, obstacles, all_tanks)
                
            if FIRE.value() == 1:
                player_tank.fire(bullets, obstacles)
            
            # 更新AI
            for enemy_tank in enemy_tanks:
                if enemy_tank.alive:
                    SimpleAI(enemy_tank, player_tank, obstacles, bullets, all_tanks)
            
            # 更新子弹
            i = 0
            while i < len(bullets):
                bullet = bullets[i]
                
                # 先清除子弹当前位置
                if (bullet[0], bullet[1]) in obstacles:
                    DrawCell(bullet[0], bullet[1], ObstacleColor)
                else:
                    DrawCell(bullet[0], bullet[1], MapColor)
                
                # 移动子弹
                if bullet[2] == 0:  # 上
                    bullet[1] -= 1
                elif bullet[2] == 1:  # 右
                    bullet[0] += 1
                elif bullet[2] == 2:  # 下
                    bullet[1] += 1
                else:  # 左
                    bullet[0] -= 1
                
                # 检查是否碰到障碍物（包括边界）- 这对所有子弹都适用
                if IsObstacle(bullet[0], bullet[1], obstacles):
                    bullets.pop(i)
                    continue
                
                # 检查玩家坦克 - 只有敌方子弹才能击中玩家
                hit_tank = False
                if player_tank.alive and player_tank.x == bullet[0] and player_tank.y == bullet[1]:
                    # 只有敌方子弹才能击中玩家
                    if not bullet[3]:  # 敌方子弹
                        player_tank.alive = False
                        bullets.pop(i)
                        hit_tank = True
                        enemy_score += 1
                    else:  # 玩家自己的子弹，忽略
                        pass
                
                # 检查敌人坦克 - 只有玩家子弹才能击中敌人，且敌方子弹不会伤害其他敌人
                if not hit_tank:
                    for j, enemy_tank in enumerate(enemy_tanks):
                        if enemy_tank.alive and enemy_tank.x == bullet[0] and enemy_tank.y == bullet[1]:
                            # 只有玩家子弹才能击中敌人
                            if bullet[3]:  # 玩家子弹
                                enemy_tank.alive = False
                                bullets.pop(i)
                                hit_tank = True
                                player_score += 1
                            break  # 跳出循环，即使有多个敌人坦克在同一位置)
                
                if hit_tank:
                    # 如果击中坦克，需要重绘被击中的坦克位置
                    if bullet[3]:  # 玩家子弹击中敌人
                        for enemy_tank in enemy_tanks:
                            if not enemy_tank.alive:
                                DrawCell(enemy_tank.x, enemy_tank.y, MapColor)
                    else:  # 敌方子弹击中玩家
                        if not player_tank.alive:
                            DrawCell(player_tank.x, player_tank.y, MapColor)
                    continue
                
                # 绘制子弹新位置
                DrawCell(bullet[0], bullet[1], BulletColor)
                i += 1
            
            # 检查子弹之间的碰撞
            CheckBulletCollision(bullets)
            
            # 更新坦克
            player_tank.update()
            for enemy_tank in enemy_tanks:
                enemy_tank.update()
                
            # 更新分数和难度显示
            DrawScore()
            DrawDifficultyLevel(len(enemy_tanks))
            
            # 检查游戏结束条件
            game_ended = False
            win = False
            
            # 检查是否所有敌人都被消灭
            all_enemies_dead = True
            for enemy_tank in enemy_tanks:
                if enemy_tank.alive:
                    all_enemies_dead = False
                    break
            
            # 如果所有敌人都被消灭，玩家获胜并提升难度
            if all_enemies_dead:
                game_ended = True
                win = True
                # 提升难度级别（最高不超过3）
                difficulty_level = min(difficulty_level + 1, 3)
            
            # 如果玩家死亡，检查分数决定胜负
            if not player_tank.alive:
                game_ended = True
                win = (player_score > enemy_score)
            
            if game_ended:
                game_state = GAME_OVER
                game_over_option = 0  # 重置选项为默认
                # 绘制游戏结束画面
                GameOver(game_over_option)
            
        elif game_state == GAME_OVER:
            # 处理上下键选择选项
            if UP.value() == 1:
                game_over_option = 0  # 选择继续游戏
                GameOver(game_over_option)
                # 等待按键释放
                while UP.value() == 1:
                    utime.sleep_ms(50)
            elif DOWN.value() == 1:
                game_over_option = 1  # 选择返回主页面
                GameOver(game_over_option)
                # 等待按键释放
                while DOWN.value() == 1:
                    utime.sleep_ms(50)
            
            # 处理发射键确认选择
            if FIRE.value() == 1:
                if game_over_option == 0:  # 继续游戏
                    game_state = PLAYING
                    # 重新创建障碍物
                    obstacles = []
                    for i in range(40):
                        x = urandom.randint(2, mapWidth - 3)
                        y = urandom.randint(2, mapHeight - 3)
                        obstacles.append((x, y))
                    
                    # 重新创建玩家坦克
                    player_tank = Tank(2, 2, 1, PlayerColor, True)
                    
                    # 重新创建敌人坦克，根据当前难度级别
                    enemy_tanks = []
                    if difficulty_level >= 1:
                        enemy_tanks.append(Tank(mapWidth - 3, mapHeight - 3, 3, EnemyColor))
                    
                    if difficulty_level >= 2:
                        enemy_tanks.append(Tank(2, mapHeight - 3, 0, EnemyColor))
                    
                    if difficulty_level >= 3:
                        enemy_tanks.append(Tank(mapWidth - 3, 2, 2, EnemyColor))
                    
                    # 清空子弹
                    bullets = []
                    prev_bullets = []
                    
                    # 重新绘制地图和界面元素
                    display.fill(MapColor)
                    DrawMap(obstacles)
                    DrawScore()
                    DrawDifficultyLevel(len(enemy_tanks))
                else:  # 返回主页面
                    game_state = START_SCREEN
                    StartScreen()
                
                # 等待按键释放
                while FIRE.value() == 1:
                    utime.sleep_ms(50)
        
        # 控制游戏速度
        utime.sleep_ms(50)

# 启动游戏
if __name__ == '__main__':
    main()